import numpy as np
from collections import Counter
from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split


class NaiveBayes:
    def __init__(self, alpha=1.0):
        """
        初始化朴素贝叶斯分类器

        参数:
            alpha (float): 拉普拉斯平滑参数，用于避免在计算概率时出现零概率问题，默认值为1.0
        """
        self.alpha = alpha
        self.class_priors = {}
        self.feature_likelihoods = {}
        self.classes = []

    def fit(self, X, y):
        """
        使用训练数据训练朴素贝叶斯分类器

        参数:
            X (numpy.ndarray): 训练数据，二维数组，每一行表示一个样本，每一列表示一个特征
            y (numpy.ndarray): 训练数据的标签，一维数组，每个元素对应一个样本的类别
        """
        # 获取训练数据中的唯一类别，存储在self.classes中
        self.classes = np.unique(y)
        num_samples, num_features = X.shape

        # 计算每个类别的先验概率 P(C)，即每个类别在训练数据中出现的频率
        self.class_priors = {c: np.mean(y == c) for c in self.classes}

        # 计算条件概率 P(X|C)，使用拉普拉斯平滑来避免零概率问题
        self.feature_likelihoods = {}
        for c in self.classes:
            # 选取属于类别c的所有样本
            X_c = X[y == c]

            likelihoods = {}
            for feature_idx in range(num_features):
                # 统计属于类别c的样本中每个特征值出现的次数
                feature_counts = Counter(X_c[:, feature_idx])

                # 计算特征的唯一值数量，加1是为了考虑拉普拉斯平滑中未观察到的值
                unique_feature_values = len(feature_counts) + 1
                # 对每个特征值的计数应用拉普拉斯平滑，即在计数上加alpha
                smoothed_counts = {feat_val: count + self.alpha for feat_val, count in feature_counts.items()}
                # 计算总计数，考虑了拉普拉斯平滑后的所有可能情况
                total_count = X_c.shape[0] + self.alpha * unique_feature_values

                # 计算每个特征值的条件概率，即平滑后的计数除以总计数
                likelihoods[feature_idx] = {feat_val: count / total_count for feat_val, count in smoothed_counts.items()}
                # 添加一个默认值，对应于未观察到的特征值，其概率为alpha / total_count
                likelihoods[feature_idx][None] = self.alpha / total_count

            self.feature_likelihoods[c] = likelihoods

    def predict(self, X):
        """
        使用训练好的模型对新数据进行预测

        参数:
            X (numpy.ndarray): 预测数据，二维数组，每一行表示一个待预测样本，每一列表示一个特征

        返回:
            numpy.ndarray: 预测结果，一维数组，每个元素对应一个样本的预测类别
        """
        predictions = []
        for sample in X:
            posteriors = {}
            for c in self.classes:
                # 初始化后验概率为类别c的先验概率的对数，使用对数概率可防止数值下溢
                posterior = np.log(self.class_priors[c])

                for feature_idx, feature_value in enumerate(sample):
                    # 获取特征值在类别c下的条件概率，如果特征值未出现过，则使用默认值（对应未观察到的特征值的概率）
                    likelihood = self.feature_likelihoods[c][feature_idx].get(feature_value, self.feature_likelihoods[c][feature_idx][None])
                    # 将条件概率的对数加到后验概率上
                    posterior += np.log(likelihood)

                posteriors[c] = posterior

            # 选择后验概率最大的类别作为预测结果
            predictions.append(max(posteriors, key=posteriors.get))

        return np.array(predictions)


# 函数功能：加载并预处理20Newsgroups数据集，将文本数据转换为特征向量，并划分为训练集和测试集
# 返回值说明：返回一个包含训练集和测试集数据和标签的元组，依次为X_train, X_test, y_train, y_test
def load_data():
    """
    加载并预处理20Newsgroups数据集

    返回:
        tuple: 包含训练集和测试集的数据和标签，依次为X_train, X_test, y_train, y_test
    """
    # 加载20Newsgroups数据集，仅选择两个类别（'sci.space'和'rec.autos'）以简化问题
    newsgroups = fetch_20newsgroups(subset='all', categories=['sci.space','rec.autos'])

    # 使用CountVectorizer将文本数据转换为特征向量，设置为二进制模式（只考虑特征是否出现），
    # 去除英语停用词，并限制特征数量为1000（选择最常见的1000个特征）
    vectorizer = CountVectorizer(binary=True, stop_words='english', max_features=1000)
    X = vectorizer.fit_transform(newsgroups.data).toarray()
    y = newsgroups.target

    # 将数据集划分为训练集和测试集，测试集占比20%，并设置随机种子为42以确保结果可复现
    return train_test_split(X, y, test_size=0.2, random_state=42)


if __name__ == "__main__":
    # 加载数据，得到训练集和测试集的特征向量和标签
    X_train, X_test, y_train, y_test = load_data()

    # 创建朴素贝叶斯模型实例，并传入拉普拉斯平滑参数alpha=1.0
    model = NaiveBayes(alpha=1.0)
    # 使用训练集数据训练模型
    model.fit(X_train, y_train)

    # 使用训练好的模型对测试集数据进行预测
    y_pred = model.predict(X_test)

    # 计算预测准确率，即预测正确的样本数占测试集样本总数的比例
    accuracy = np.mean(y_pred == y_test)
    print(f"测试集准确率: {accuracy:.4f}")